import numpy as np
import torch
import joblib
import random
import os
from deepctr_torch.inputs import SparseFeat, DenseFeat, VarLenSparseFeat, get_feature_names
from deepctr_torch.models import DIEN

MAX_LEN = 50

def seed_torch(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
seed_torch(1234)


feature_columns = [SparseFeat('customer_id', 1371980, embedding_dim=64, use_hash=False),
                   SparseFeat('FN', 2, embedding_dim=4, use_hash=False),
                   SparseFeat('Active', 2, embedding_dim=4, use_hash=False),
                   SparseFeat('club_member_status', 2, embedding_dim=4, use_hash=False),
                   SparseFeat('fashion_news_frequency', 2, embedding_dim=4, use_hash=False),
                   SparseFeat('artid', 105542 + 1, embedding_dim=64, use_hash=False),
                   SparseFeat('pcode', 47224 + 1, embedding_dim=64, use_hash=False),
                   DenseFeat('age', 1)]
feature_columns +=[VarLenSparseFeat(SparseFeat('hist_artid', vocabulary_size=105542 + 1, embedding_dim=64, embedding_name='artid'), maxlen=MAX_LEN, length_name="seq_length"),
                   VarLenSparseFeat(SparseFeat('hist_pcode', vocabulary_size=47224 + 1, embedding_dim=64, embedding_name='pcode'), maxlen=MAX_LEN,length_name="seq_length")]
behavior_feature_list = ["artid", "pcode"]

x_train = joblib.load("../Data/Processed/train_dict.db")
x_valid = joblib.load("../Data/Processed/valid_dict.db")
y_train = x_train['y']
y_valid = x_valid['y']
x_train = {name: x_train[name] for name in get_feature_names(feature_columns)}
x_valid = {name: x_valid[name] for name in get_feature_names(feature_columns)}

device = 'cpu'
use_cuda = True
if use_cuda and torch.cuda.is_available():
    print('cuda ready...')
    device = 'cuda:0'


model = DIEN(feature_columns, behavior_feature_list,
             gru_type="AUGRU", use_negsampling=False, device=device)

model.compile(torch.optim.Adam(model.parameters(), 0.0005), 'binary_crossentropy',
              metrics=['binary_crossentropy', 'auc'])
    
history = model.fit(x_train, y_train, batch_size=1024, epochs=10, verbose=1, validation_data=(x_valid, y_valid), shuffle=False)
